#!/usr/local/bin/python
# Retrieves from BLAST results top neighbors that are cultivated.
# A little rough right now.. because using "strain" as a keyword to filter out uncultivated neighbors.

import sys
import re
from Bio import SeqIO
from Bio import Entrez

Entrez.email = "jimmysaw@gmail.com"

filename = sys.argv[1]
f = open(filename, "rU")
fh = f.readlines()

max2retrieve = int(sys.argv[2])

m1 = re.compile('Unculture.*', re.IGNORECASE)
m2 = re.compile('.*\|(.*)\|')
m3 = re.compile('complete genome', re.IGNORECASE)
m4 = re.compile('Submitted', re.IGNORECASE)
m5 = re.compile('[K|M|R|Y|S|W|B|V|H|D|X|N]')
m6 = re.compile('Published Only in Database', re.IGNORECASE)
m7 = re.compile('Int\. J\. Syst\. Evol\. Microbiol\.')
m8 = re.compile('strain')

seqid = ""

uncl_acc = []
uncl_def = []
uncl_sims = {}

compl_acc = []
compl_def = []
compl_sims = {}

hit_acc = []
hit_def = []
hit_sims = {}

for line in fh:
    l = line.split('\t')
    seqid = l[5]
#    if m1.match(l[13].rstrip()):
#        pattern = m2.match(l[9])
#        uncl_acc.append(pattern.group(1))
#        uncl_sims[pattern.group(1)] = l[3]
#        uncl_def.append(l[13].rstrip())
#    elif m3.search(l[13].rstrip()):
#        pattern = m2.match(l[9])
#        compl_acc.append(pattern.group(1))
#        compl_sims[pattern.group(1)] = l[3]
#        compl_def.append(l[13].rstrip())
#    else:
#        pattern = m2.match(l[9])
#        hit_acc.append(pattern.group(1))
#        hit_sims[pattern.group(1)] = l[3]
#        hit_def.append(l[13].rstrip())
    if m2.match(l[9]):
        pattern = m2.match(l[9])
        hit_acc.append(pattern.group(1))
        hit_sims[pattern.group(1)] = l[3]
        hit_def.append(l[13].rstrip())

#print "## Sequences from complete genomes: ", len(compl_acc)
#print "## Sequences from uncultivated organisms: ", len(uncl_acc)
#print "## Sequences from other organisms: ", len(hit_acc), "\n"

unique_hit_acc = set(hit_acc)
unique_list = list(unique_hit_acc)

max_num = 0

if len(unique_list) > 500:
    max_num = 500
else:
    max_num = len(unique_list)

top_hits = []
top_seqs = []
top_sims = []
top_journal = []
top_titles = []

for acc in unique_list[0:max_num]:
    handle = Entrez.efetch(db="nucleotide", id=acc, rettype="gb")
    rec = SeqIO.read(handle, "genbank")
    journal_list = []
    ref_title = []
    mol_list = []
    for ref in rec.annotations['references']:
        journal_list.append(ref.journal)
        ref_title.append(ref.title)
    for feat in rec.features:
        if feat.qualifiers.has_key('mol_type'):
            x = feat.qualifiers.get('mol_type')
            mol_list.append(x[0])

    if m5.search(rec.seq.tostring()):
        pass
    elif len(rec.seq) < 1000:
        pass
    else:
        top_hits.append(acc)
        top_seqs.append(rec)
        top_sims.append(hit_sims[acc])
        top_journal.append(journal_list[0])
        top_titles.append(ref_title[0])

#print "## Top " + str(len(top_hits)) + " sequences (>= 1000bp): ", "\n"

ijsem_hits = []
ijsem_seqs = []
ijsem_sims = []
ijsem_jour = []
ijsem_title = []

other_hits = []
other_seqs = []
other_sims = []
other_jour = []
other_title = []

for i, j in enumerate(top_hits):
    if m7.search(top_journal[i]):
        ijsem_hits.append(top_hits[i])
        ijsem_seqs.append(top_seqs[i])
        ijsem_sims.append(top_sims[i])
        ijsem_jour.append(top_journal[i])
        ijsem_title.append(top_titles[i])
    else:
        other_hits.append(top_hits[i])
        other_seqs.append(top_seqs[i])
        other_sims.append(top_sims[i])
        other_jour.append(top_journal[i])
        other_title.append(top_titles[i])

#if len(ijsem_hits) > 0:
#    print "## IJSEM hits: \n"
#else:
#    print "## No IJSEM hits!\n"

#for i, j in enumerate(ijsem_hits[0:max2retrieve]):
#    print str(i+1).zfill(2), ijsem_hits[i], "(" + ijsem_sims[i] + ")"
#    print str(i+1).zfill(2), ijsem_jour[i]
#    print str(i+1).zfill(2), ijsem_title[i]
#    print ""

#for seq in ijsem_seqs[0:max2retrieve]:
#    print seq.format('fasta').strip()

#print ""

#if len(other_hits) > 0:
#    print "## Other journals: \n"
#else:
#    print "## No journal hits!\n"

#for i, j in enumerate(other_hits[0:max2retrieve]):
#    print str(i+1).zfill(2), other_hits[i], "(" + other_sims[i] + ")"
#    print str(i+1).zfill(2), other_jour[i]
#    print str(i+1).zfill(2), other_title[i]
#    print ""
#    #print other_seqs[i].format('fasta').strip()

#for i, j in enumerate(top_hits[0:10]):
#    print str(i+1).zfill(2), top_hits[i], "(" + top_sims[i] + ")"
#    print str(i+1).zfill(2), top_journal[i]
#    print str(i+1).zfill(2), top_titles[i]
#    print ""
#    print top_seqs[i].format('fasta').strip()

for seq in top_seqs[0:max2retrieve]:
    print seq.format('fasta').strip()

#for seq in other_seqs[0:max2retrieve]:
#    print seq.format('fasta').strip()

print ""

f.close()
